import pygame
import random

# 初始化Pygame
pygame.init()

# 设置窗口大小
screen_width = 300
screen_height = 600
block_size = 30
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('俄罗斯方块')

# 定义颜色
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
CYAN = (0, 255, 255)
MAGENTA = (255, 0, 255)
YELLOW = (255, 255, 0)
ORANGE = (255, 165, 0)

colors = [BLACK, RED, GREEN, BLUE, CYAN, MAGENTA, YELLOW, ORANGE]
shapes = [
    [[1, 1, 1, 1]],  # I
    [[1, 1], [1, 1]],  # O
    [[0, 1, 0], [1, 1, 1]],  # T
    [[1, 0, 0], [1, 1, 1]],  # L
    [[0, 0, 1], [1, 1, 1]],  # J
    [[0, 1, 1], [1, 1, 0]],  # S
    [[1, 1, 0], [0, 1, 1]]   # Z
]

def create_shape():
    shape = random.choice(shapes)
    color = random.choice(colors[1:])
    return {'shape': shape, 'color': color, 'x': 5, 'y': 0}

def draw_grid():
    for x in range(0, screen_width, block_size):
        for y in range(0, screen_height, block_size):
            pygame.draw.rect(screen, WHITE, (x, y, block_size, block_size), 1)

def draw_shape(shape):
    for y, row in enumerate(shape['shape']):
        for x, cell in enumerate(row):
            if cell:
                pygame.draw.rect(screen, shape['color'], (shape['x'] + x * block_size, shape['y'] + y * block_size, block_size, block_size))

def check_collision(shape, grid):
    for y, row in enumerate(shape['shape']):
        for x, cell in enumerate(row):
            if cell:
                if shape['x'] + x < 0 or shape['x'] + x >= screen_width // block_size or shape['y'] + y >= screen_height // block_size:
                    return True
                if grid[shape['y'] + y][shape['x'] + x]:
                    return True
    return False

def rotate_shape(shape):
    new_shape = list(zip(*shape['shape'][::-1]))
    if not check_collision({'shape': new_shape, 'color': shape['color'], 'x': shape['x'], 'y': shape['y']}, grid):
        shape['shape'] = new_shape

def main():
    clock = pygame.time.Clock()
    grid = [[0 for _ in range(screen_width // block_size)] for _ in range(screen_height // block_size)]
    current_shape = create_shape()
    fall_time = 0

    while True:
        screen.fill(BLACK)
        draw_grid()
        draw_shape(current_shape)

        fall_speed = 0.27
        fall_time += clock.get_rawtime()
        clock.tick()

        if fall_time / 1000 >= fall_speed:
            fall_time = 0
            current_shape['y'] += 1
            if check_collision(current_shape, grid):
                current_shape['y'] -= 1
                for y, row in enumerate(current_shape['shape']):
                    for x, cell in enumerate(row):
                        if cell:
                            grid[current_shape['y'] + y][current_shape['x'] + x] = current_shape['color']
                current_shape = create_shape()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                return
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_LEFT:
                    current_shape['x'] -= 1
                    if check_collision(current_shape, grid):
                        current_shape['x'] += 1
                if event.key == pygame.K_RIGHT:
                    current_shape['x'] += 1
                    if check_collision(current_shape, grid):
                        current_shape['x'] -= 1
                if event.key == pygame.K_DOWN:
                    current_shape['y'] += 1
                    if check_collision(current_shape, grid):
                        current_shape['y'] -= 1
                if event.key == pygame.K_UP:
                    rotate_shape(current_shape)

        pygame.display.flip()

if __name__ == '__main__':
    main()
